package com.pandy.code;

import java.util.HashMap;

/**
 * @description: 前缀树
 * @author: panhong
 * @date: 2025/3/26
 */
public class Code09_Trie2 {

    /**
     * 前缀树节点
     */
    static class Node {
        public int pass;
        public int end;
        public HashMap<Integer,Node> nexts; //子节点特别多的情况

        public Node() {
            pass = 0;
            end = 0;
            nexts = new HashMap<>();
        }
    }

    /**
     * 前缀树操作
     * 通用类型
     */
    static class Trie {
        Node root;

        public Trie() {
            root = new Node();
        }

        /**
         * 插入
         * @param word
         */
        public void insert(String word) {
            if(word == null || word.isEmpty()) {
                return;
            }
            Node node = root;
            node.pass++;
            int ascii = 0;
            char[] charArray = word.toCharArray();
            for(int i=0;i<charArray.length;i++) {
                ascii = charArray[i];
                if(!node.nexts.containsKey(ascii)) {
                    node.nexts.put(ascii,new Node());
                }
                node = node.nexts.get(ascii);
                node.pass++;
            }
            node.end++;
        }

        /**
         * 删除
         * @param word
         */
        public void delete(String word) {
            int count = search(word);

            if(count > 0) {
                Node node = root;
                node.pass--;
                int ascii = 0;
                char[] charArray = word.toCharArray();
                for(int i=0;i<charArray.length;i++) {
                    ascii = charArray[i];
                    node = node.nexts.get(ascii);

                    if(node.pass-- == 0) {
                        node.nexts.remove(ascii);
                    }
                }
                node.end--;
            }

        }

        /**
         * 统计出word出现的次数
         * @param word
         * @return
         */
        public int search(String word) {
            if(word == null || word.isEmpty()) {
                return 0;
            }

            Node node = root;
            int ascii = 0;
            char[] charArray = word.toCharArray();
            for(int i=0;i<charArray.length;i++) {
                ascii = charArray[i];
                if(!node.nexts.containsKey(ascii)) {
                    return 0;
                }
                node = node.nexts.get(ascii);
            }
            return node.end;
        }

        /**
         * 统计以prefix为前缀的单词数量
         * @param prefix
         * @return
         */
        public  int prefixNumber(String prefix) {
            if(prefix == null || prefix.isEmpty()) {
                return 0;
            }

            Node node = root;
            int ascii = 0;
            char[] charArray = prefix.toCharArray();
            for(int i=0;i<charArray.length;i++) {
                ascii = charArray[i];
                if(!node.nexts.containsKey(ascii)) {
                    return 0;
                }
                node = node.nexts.get(ascii);
            }
            return node.pass;
        }
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("abc");
        trie.insert("abd");
        trie.insert("abcd");
        trie.insert("bce");
        trie.insert("bcde");
        trie.insert("aabc");
        trie.insert("abc");

        System.out.println(trie.search("aabc"));
        System.out.println(trie.prefixNumber("a"));
        System.out.println("===删除后===");
        trie.delete("aabc");
        System.out.println(trie.search("aabc"));
        System.out.println(trie.prefixNumber("a"));
    }
}
